#include <cuda_runtime.h>
#include <cuda_runtime.h>
#include <device_launch_parameters.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <time.h>
#include <sys/time.h>

const int TILE_WIDTH = 16;

double get_cpu_time_ms() {
    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC, &ts);
    return (ts.tv_sec * 1000.0) + (ts.tv_nsec / 1000000.0);
}

void MatVecMultiKernelCPU(const double* A, const double* x, double* y, int m, int n) 
{
	for (int i = 0; i < m; ++i) {
		double s = 0.0;
		for (int j = 0; j < n; ++j) {
			s += A[i * n + j] * x[j];
		}
		y[i] = s;
	}
}

__global__ void MatVecMultiKernel(const double* A, const double* x, double* y, int m, int n, bool useShared) 
{
	__shared__ double s_x[TILE_WIDTH];
    int bx = blockIdx.x;
    int tx = threadIdx.x;
	int row_idx = bx * TILE_WIDTH + tx;
	double pValue = 0.0;

    if (useShared == false) {
        for (int j = 0; j < n; j++) {
            pValue += A[row_idx * n + j] * x[j];
        }
        y[row_idx] = pValue;
        return;
    } else {
        for (int t = 0; t < (n + TILE_WIDTH - 1) / TILE_WIDTH; t++) {
            if(t * TILE_WIDTH + tx < n) {
                s_x[tx] = x[t * TILE_WIDTH + tx];
            } 
            else {
                s_x[tx] = 0.0;
            }
            __syncthreads();

            for(int i = 0; i < TILE_WIDTH; i++) {
                if(t * TILE_WIDTH + i < n) {
                    pValue += A[row_idx * n + t * TILE_WIDTH + i] * s_x[i];
                }
            }
            __syncthreads();
        }
	    y[row_idx] = pValue;
    }	
}

void verify(const double* y_cpu, const double* y, int rows) {
	double err_cnt = 0.0;
    double error_threshold = 1e-9;
    for (int r = 0; r < rows; ++r) {
        double e = fabs(y_cpu[r] - y[r]);
        if (e > error_threshold) err_cnt ++;
    }
    if (err_cnt > 0) {
        printf("Verification failed! Number of errors: %f\n", err_cnt);
    } else {
        printf("Verification passed!\n");
    }
}

int main() {
    bool useShared = false;
	const char* inputPath = "input1.txt";
	const char* outputPath = "output1.txt";
    int m = 0, n = 0;
	FILE* fin = fopen(inputPath, "r");
	fscanf(fin, "%d,%d", &m, &n);
	fclose(fin);

	size_t sizeA = (size_t)m * (size_t)n;
    size_t sizex = (size_t)n;
    size_t sizey = (size_t)m;
	double* h_A = (double*)malloc(sizeA * sizeof(double));
	double* x = (double*)malloc(sizex * sizeof(double));
	double* h_y_cpu = (double*)malloc(sizey * sizeof(double));
	double* h_y = (double*)malloc(sizey * sizeof(double));

	for (size_t i = 0; i < sizeA; ++i) h_A[i] = (double)rand() / (double)RAND_MAX;
	for (int i = 0; i < n; ++i) x[i] = (double)rand() / (double)RAND_MAX;

	double cpu_start, cpu_stop, cpu_time;
    cpu_start = get_cpu_time_ms();
	MatVecMultiKernelCPU(h_A, x, h_y_cpu, m, n);
	cpu_stop = get_cpu_time_ms();
    cpu_time = cpu_stop - cpu_start;
    printf("CPU Time: %.2f ms\n", cpu_time);

	double *d_A = NULL, *d_x = NULL, *d_y = NULL;
	cudaMalloc((void**)&d_A, sizeA * sizeof(double));
	cudaMalloc((void**)&d_x, (size_t)n * sizeof(double));
	cudaMalloc((void**)&d_y, (size_t)m * sizeof(double));
	cudaMemcpy(d_A, h_A, sizeA * sizeof(double), cudaMemcpyHostToDevice);
	cudaMemcpy(d_x, x, (size_t)n * sizeof(double), cudaMemcpyHostToDevice);

    cudaEvent_t gpuStart, gpuStop;
    float ElapsedTime;
    cudaEventCreate(&gpuStart);
    cudaEventCreate(&gpuStop);
    cudaEventRecord(gpuStart);
    dim3 grid((m + TILE_WIDTH - 1) / TILE_WIDTH);
    dim3 block(TILE_WIDTH);

	MatVecMultiKernel<<<grid, block>>>(d_A, d_x, d_y, m, n, useShared);

	cudaEventRecord(gpuStop);
    cudaEventSynchronize(gpuStop);
    cudaEventElapsedTime(&ElapsedTime, gpuStart, gpuStop);
    double gpu_time = (double)ElapsedTime;
    printf("GPU Time: %.2f ms\n", gpu_time);

	cudaMemcpy(h_y, d_y, (size_t)m * sizeof(double), cudaMemcpyDeviceToHost);

	verify(h_y_cpu, h_y, m);

	FILE* fout = fopen(outputPath, "w");
	if (fout) {
		// write CPU and GPU times in ms with two decimals
		fprintf(fout, "%.2f,%.2f", cpu_time, gpu_time);
		fclose(fout);
	} else {
		fprintf(stderr, "Failed to open %s for writing\n", outputPath);
	}

	// cleanup
	free(h_A);
	free(x);
	free(h_y_cpu);
	free(h_y);
	cudaFree(d_A);
	cudaFree(d_x);
	cudaFree(d_y);

	return 0;
}



